"""
核心环境模块
包含地图定义、物理规则、锁定判断等基础环境功能，不包含强化学习相关逻辑

作者：wjttdbx
版本：3.0 - 重构版本，使用工具模块
"""
import sys
sys.path.insert(0, sys.path[0]+"/../")

import numpy as np
import math
from typing import List, Tuple, Optional
from .config import GameConfig
from .units import Unit, WhiteBoat, WhiteDrone, BlackBoat
from .LockingSystem import LockingSystem
from .CollisionSystem import CollisionSystem
from .DroneManager import DroneManager
from .BoundarySystem import BoundarySystem
from .utils import (
    DEPLOYMENT_LINE_START, DEPLOYMENT_LINE_END, TRAPEZOID_POINTS,
    is_in_black_init_area, get_formation_position
)
   


class CoreEnvironment:
    """
    核心环境类
    
    管理游戏的基础环境，包括地图、物理规则、单位初始化等
    不包含强化学习相关逻辑
    """
    
    def __init__(self):
        """初始化核心环境"""
        self.debug_mode = False  # 是否启用调试模式
        self.white_boats: List[WhiteBoat] = []
        self.all_white_drones: List[WhiteDrone] = [] # 统一管理所有无人机
        self.black_boats: List[BlackBoat] = []
        
        # 使用新的无人机管理器
        self.drone_manager = DroneManager()
        
        # 统计数据
        self.collision_count = 0
        # 新口径（采样）统计：仅统计涉及白方的碰撞，并按10秒采样累计
        self.white_collision_count_sampled = 0
        self.next_collision_sample_time = 10.0
        self.white_lock_count = 0
        self.black_lock_count = 0
        self.intercepted_black_boats = 0
        # 第一次黑方突破的时间（从游戏开始计时的秒数）。无突破则为None
        self.first_breakthrough_time = None
        
        # 时间追踪
        self.current_time = 0.0

    def reset(self):
        """重置环境到初始状态"""
        self.current_time = 0.0
        self.collision_count = 0
        self.white_collision_count_sampled = 0
        self.next_collision_sample_time = 10.0
        self.white_lock_count = 0
        self.black_lock_count = 0
        self.intercepted_black_boats = 0
        self.first_breakthrough_time = None
        
        # 初始化单位
        self._initialize_units()
        
    def _initialize_units(self):
        """初始化所有单位，包括白方和黑方"""
        self.white_boats = []
        self.all_white_drones = []
        self.black_boats = []
        
        # 初始化白方单位
        num_white = GameConfig.WHITE_BOATS
        center_x = DEPLOYMENT_LINE_START[0]
        center_y = (DEPLOYMENT_LINE_START[1] + DEPLOYMENT_LINE_END[1]) / 2
        spacing = 5000
        
        for i in range(num_white):
            x, y = get_formation_position(i, num_white, center_x, center_y, spacing)
            boat = WhiteBoat(i, x, y)
            self.white_boats.append(boat)

            # 每艘无人艇初始携带1架满能无人机
            # 无人机ID与母舰ID关联，但可以更灵活，例如使用 100+i
            drone_id = i 
            drone = WhiteDrone(drone_id, x, y, parent_boat_id=i)
            drone.energy = 100.0
            drone.state = "recharging"  # 初始状态为充能中
            # drone.docked_on_boat_id 已经在其 __init__ 中被设置为 parent_boat_id
            self.all_white_drones.append(drone)
        
        # 初始化黑方单位
        self._initialize_black_boats()

    def _initialize_white_boats(self):
        """此方法已合并到 _initialize_units，保留以防有旧代码依赖"""
        pass
    
    def _initialize_black_boats(self):
        """初始化黑方无人艇在A2A3A4A5梯形区域"""
        num_black = GameConfig.BLACK_BOATS

        # 获取梯形区域的边界
        min_x = min(p[0] for p in TRAPEZOID_POINTS)
        max_x = max(p[0] for p in TRAPEZOID_POINTS)
        min_y = min(p[1] for p in TRAPEZOID_POINTS)
        max_y = max(p[1] for p in TRAPEZOID_POINTS)
        
        for i in range(num_black):
            attempts = 0
            while attempts < 100:  # 防止无限循环
                x = np.random.uniform(min_x, max_x)
                y = np.random.uniform(min_y, max_y)
                
                if is_in_black_init_area(x, y):
                    boat = BlackBoat(i, x, y)
                    self.black_boats.append(boat)
                    break
                attempts += 1
    
    def update(self, dt: float):
        """
        更新环境状态
        
        Args:
            dt (float): 时间步长
        """
        self.current_time += dt
        
        # 更新锁定系统
        #self._update_locking_system()
        
        # 使用新的无人机管理器更新无人机系统
        self.drone_manager.update_all_drones(self.white_boats, self.all_white_drones, self.current_time, dt)

        # 检查无人船碰撞
        collision_increment = CollisionSystem.check_all_collisions(
            self.white_boats, self.black_boats)
        self.collision_count += collision_increment
        
        # 边界约束
        self._apply_boundary_constraints()
        
        # 检查突破
        self._check_breakthroughs()

        # 每10秒采样一次，仅统计涉及白方的碰撞，并累计到新口径计数中
        while self.current_time >= self.next_collision_sample_time - 1e-6:
            sampled_now = CollisionSystem.count_white_involved_collisions(
                self.white_boats, self.black_boats
            )
            self.white_collision_count_sampled += sampled_now
            self.next_collision_sample_time += 10.0
    
    def _update_locking_system(self):
        """更新锁定系统"""
        # 获取飞行中的无人机
        flying_drones = [d for d in self.all_white_drones if d.state == 'flying']
        
        # 更新白方无人机探测系统
        for drone in flying_drones:
            LockingSystem.update_white_drone_detection(drone, self.black_boats)
        
        # 更新白方无人船探测系统（包含跨平台锁定接力）
        for boat in self.white_boats:
            if boat.active and not boat.is_frozen(self.current_time):
                # 使用支持跨平台锁定的探测更新方法
                LockingSystem.update_white_boat_detection_legacy(boat, self.black_boats, self.all_white_drones, self.current_time)
        
        # 更新黑方锁定系统
        for boat in self.black_boats:
            if boat.active and not boat.is_frozen(self.current_time):
                LockingSystem.update_black_boat_detection(boat, self.white_boats)
    
    def _apply_boundary_constraints(self):
        """应用边界约束"""
        for boat in self.white_boats + self.black_boats:
            BoundarySystem.enforce_boundaries(boat)
        
        for drone in self.all_white_drones:
            # 只对飞行中的无人机应用边界约束
            if drone.state == 'flying':
                BoundarySystem.enforce_boundaries(drone)
    
    def _check_breakthroughs(self):
        """检查突破情况"""
        for black_boat in self.black_boats:
            if black_boat.active and not black_boat.crossed:
                BoundarySystem.check_breakthrough(black_boat)
                # 记录第一次突破时间
                if black_boat.crossed and self.first_breakthrough_time is None:
                    self.first_breakthrough_time = self.current_time
    
    def try_white_boat_lock(self, boat_id: int) -> bool:
        """
        尝试白方无人艇锁定最近的目标
        
        Args:
            boat_id (int): 无人艇ID
            
        Returns:
            bool: 是否成功锁定
        """
        if boat_id < len(self.white_boats):
            boat = self.white_boats[boat_id]
            if boat.active and not boat.is_frozen(self.current_time):
                # 自动选择最近的目标
                if boat.detected_targets:
                    nearest_target = min(boat.detected_targets, key=lambda t: boat.distance_to(t))
                    success = LockingSystem.try_white_boat_lock(boat, nearest_target.id, self.current_time)
                    if success:
                        self.white_lock_count += 1
                        # 检查是否拦截成功
                        if boat.current_target and not boat.current_target.active:
                            self.intercepted_black_boats += 1
                    return success
        return False
    
    def try_white_boat_lock_target(self, boat_id: int, target_id: int) -> bool:
        """
        尝试白方无人艇锁定指定目标
        
        锁定判定逻辑（与无人机指引一致）：
        1. 初次锁定：目标已被探测（可由无人机共享进入 detected_targets），且距离≤40km
        2. 持续锁定：目标只需在锁定范围内（≤40km）
        3. 目标切换：新目标必须在探测范围内（≤20km）
        
        Args:
            boat_id (int): 无人艇ID
            target_id (int): 目标ID

        Returns:
            bool: 是否成功锁定
        """
        if GameConfig.debug_mode:
            print(f"[CoreEnvironment] 尝试无人艇{boat_id} 锁定目标索引{target_id}")

        if boat_id < len(self.white_boats):
            boat = self.white_boats[boat_id]
            if boat.active and not boat.is_frozen(self.current_time):
                if GameConfig.debug_mode:
                    print(f"[CoreEnvironment] 无人艇{boat_id} 状态正常")

                # 检查目标是否在有效范围内
                aim_target = None
                
                for idx, target in enumerate(boat.detected_targets):
                    if target.id == target_id:
                        aim_target = target
                        break

                if aim_target is not None and aim_target.active:
                    distance = boat.distance_to(aim_target)
                    
                    # 判定锁定类型和范围要求
                    lock_type = "未知"
                    range_required = 0
                    range_check_passed = False
                    
                    if not boat.is_locking:
                        # 初次锁定：目标已被探测（含无人机共享），允许在40km内直接开始锁定
                        lock_type = "初次锁定"
                        range_required = GameConfig.WHITE_BOAT_LOCK_RANGE
                        range_check_passed = distance <= GameConfig.WHITE_BOAT_LOCK_RANGE
                        if GameConfig.debug_mode:
                            print(f"[CoreEnvironment] 无人艇{boat_id} 初次锁定目标{aim_target.id}: 距离={distance:.0f}m, 需要锁定范围≤{range_required}m")
                    elif boat.current_target and boat.current_target.id == target_id:
                        # 持续锁定：只需要锁定范围
                        lock_type = "持续锁定"
                        range_required = GameConfig.WHITE_BOAT_LOCK_RANGE
                        range_check_passed = distance <= GameConfig.WHITE_BOAT_LOCK_RANGE
                        if GameConfig.debug_mode:
                            print(f"[CoreEnvironment] 无人艇{boat_id} 持续锁定目标{aim_target.id}: 距离={distance:.0f}m, 需要锁定范围≤{range_required}m")
                    else:
                        # 目标切换：需要探测范围
                        lock_type = "目标切换"
                        range_required = GameConfig.WHITE_BOAT_DETECT_RANGE
                        range_check_passed = distance <= GameConfig.WHITE_BOAT_DETECT_RANGE
                        if GameConfig.debug_mode:
                            print(f"[CoreEnvironment] 无人艇{boat_id} 切换锁定目标{aim_target.id}: 距离={distance:.0f}m, 需要探测范围≤{range_required}m")
                    
                    if range_check_passed:
                        if GameConfig.debug_mode:
                            print(f"[CoreEnvironment] 无人艇{boat_id} {lock_type}目标{aim_target.id} 范围检查通过: {distance:.0f}m ≤ {range_required}m")

                        boat.start_locking(aim_target, self.current_time)
                        if GameConfig.debug_mode:
                            lock_duration = boat.get_lock_duration(self.current_time)
                            print(f"[CoreEnvironment] 无人艇{boat_id} 开始锁定目标{aim_target.id}, 当前锁定时间: {lock_duration:.1f}s")

                        # 检查是否锁定成功
                        can_succeed = LockingSystem.can_lock_successfully(
                            boat,
                            self.current_time,
                            GameConfig.WHITE_BOAT_LOCK_DURATION,
                            GameConfig.WHITE_BOAT_LOCK_SUCCESS_RATE
                        )

                        if can_succeed:
                            # 锁定成功处理
                            print(f"[CoreEnvironment] 无人艇{boat_id} 成功锁定目标 {aim_target.id}")
                            LockingSystem.apply_lock_success(aim_target, self.current_time)
                            boat.stop_locking()
                            self.white_lock_count += 1
                            # 检查是否拦截成功
                            if not aim_target.active:
                                self.intercepted_black_boats += 1
                                if GameConfig.debug_mode:
                                    print(f" [CoreEnvironment] 目标{aim_target.id} 被消灭")
                            return True
                        return True  # 正在锁定中
                    else:
                        if GameConfig.debug_mode:
                            print(f" [CoreEnvironment] 无人艇{boat_id} {lock_type}目标{aim_target.id} 范围检查失败: {distance:.0f}m > {range_required}m")
                        return False
                else:
                    if GameConfig.debug_mode:
                        if aim_target is not None:
                            print(f" [CoreEnvironment] 无人艇{boat_id} 目标 {aim_target.id}已失效")
                        else:
                            print(f" [CoreEnvironment] 无人艇{boat_id} 目标索引 {target_id} 超出范围，探测到 {len(boat.detected_targets)} 个目标")
                    return False
            else:
                if GameConfig.debug_mode:
                    status = "已失效" if not boat.active else "被冻结"
                    print(f" [CoreEnvironment] 无人艇{boat_id} {status}")
        else:
            if GameConfig.debug_mode:
                print(f" [CoreEnvironment] 无人艇ID {boat_id} 超出范围，总共有 {len(self.white_boats)} 艘无人艇")
        return False
    
    def stop_white_boat_lock(self, boat_id: int):
        """
        停止白方无人艇的锁定
        
        Args:
            boat_id (int): 无人艇ID
        """
        if boat_id < len(self.white_boats):
            boat = self.white_boats[boat_id]
            if boat.active:
                boat.stop_locking()
                if self.debug_mode:
                    print(f"[CoreEnvironment] 无人艇{boat.id} 停止锁定")
    
    def try_black_boat_lock(self, boat_id: int) -> bool:
        """
        尝试黑方无人艇锁定
        
        Args:
            boat_id (int): 无人艇ID
            
        Returns:
            bool: 是否成功锁定
        """
        if boat_id < len(self.black_boats):
            boat = self.black_boats[boat_id]
            if boat.active and not boat.is_frozen(self.current_time):
                success = LockingSystem.try_black_boat_lock(boat, self.current_time)
                if success:
                    self.black_lock_count += 1
                return success
        return False
    
    def update_detection_system(self):
        """
        更新探测系统
        
        提供给外部调用的公共接口，用于更新所有单位的探测状态。
        这确保AI能够基于最新的探测信息做出决策。
        """
        self._update_locking_system()
    
    def launch_drone(self, boat_id: int) -> bool:
        """
        发射无人机
        
        Args:
            boat_id (int): 无人艇ID
            
        Returns:
            bool: 是否成功发射
        """
        # 检查母舰ID是否有效
        if boat_id >= len(self.white_boats):
            print(f"[CoreEnvironment] 母舰{boat_id}不存在")
            return False
        
        boat = self.white_boats[boat_id]
        
        # 检查母舰状态
        if not boat.active:
            print(f"[CoreEnvironment] 母舰{boat_id}已被摧毁，无法发射无人机")
            return False
        
        if boat.is_frozen(self.current_time):
            print(f"[CoreEnvironment] 母舰{boat_id}处于冻结状态，无法发射无人机")
            return False
        
        # drone_manager 现在需要知道所有无人机的状态来做决策
        success = self.drone_manager.launch_drone(boat, self.all_white_drones, self.current_time)
        
        if not success and self.debug_mode:
            print(f"[CoreEnvironment] 母舰{boat_id}发射无人机失败")
        
        return success
    
    def land_drone(self, drone_id: int, target_boat_id: int) -> bool:
        """
        无人机降落
        
        Args:
            drone_id (int): 无人机ID
            target_boat_id (int): 目标母舰ID
            
        Returns:
            bool: 是否成功降落
        """
        # 找到指定的无人机
        drone = None
        for d in self.all_white_drones:
            if d.id == drone_id:
                drone = d
                break
        
        if not drone:
            print(f"[CoreEnvironment] 未找到无人机{drone_id}")
            return False
        
        # 找到指定的母舰
        if target_boat_id >= len(self.white_boats):
            print(f"[CoreEnvironment] 目标母舰{target_boat_id}不存在")
            return False
        
        target_boat = self.white_boats[target_boat_id]
        
        # 调用DroneManager执行降落
        return self.drone_manager.land_drone(target_boat, drone, self.current_time)
    
    def get_game_state(self) -> dict:
        """
        获取游戏状态信息
        
        Returns:
            dict: 游戏状态字典
        """
        active_white = sum(1 for boat in self.white_boats if boat.active)
        active_black = sum(1 for boat in self.black_boats if boat.active)
        crossed_black = sum(1 for boat in self.black_boats if boat.crossed)
        eliminated_black = sum(1 for boat in self.black_boats if not boat.active)
        
        # 获取无人机状态统计
        drone_status = self.drone_manager.get_drone_status(self.all_white_drones)
        
        return {
            'current_time': self.current_time,
            'active_white_boats': active_white,
            'active_black_boats': active_black,
            'crossed_black_boats': crossed_black,
            'eliminated_black_boats': eliminated_black,
            'total_black_boats': len(self.black_boats),
            # 对外暴露的新口径：每10秒采样一次、仅涉及白方的碰撞累计
            'collision_count': self.white_collision_count_sampled,
            # 可选：保留原始逐帧累计供诊断（不在UI中使用）
            'collision_count_raw': self.collision_count,
            'white_lock_count': self.white_lock_count,
            'black_lock_count': self.black_lock_count,
            'intercepted_black_boats': self.intercepted_black_boats,
            'first_breakthrough_time': self.first_breakthrough_time,
            'flying_drones': sum(1 for d in self.all_white_drones if d.state == 'flying'),
            'drone_status': drone_status,
        }
    
    def is_game_over(self) -> Tuple[bool, str]:
        """
        检查游戏是否结束
        
        Returns:
            Tuple[bool, str]: (是否结束, 结束原因)
        """
        all_black = sum(1 for boat in self.black_boats)
        active_black = sum(1 for boat in self.black_boats if boat.active)
        crossed_black = sum(1 for boat in self.black_boats if boat.crossed)
        active_white = sum(1 for boat in self.white_boats if boat.active)
        
        # 条件1：所有黑方无人艇都被击沉或消失（active_black == 0）
        if active_black == 0:
            return True, "white_wins"  # 白方胜利

        # 条件2：白方全灭
        if active_white == 0:
            return True, "black_wins"  # 白方全灭，游戏结束
        
        return False, "ongoing"
    
    def get_environment_info(self) -> dict:
        """
        获取环境详细信息
        
        Returns:
            dict: 环境信息字典
        """
        return {
            'white_boats': len(self.white_boats),
            'white_drones': len(self.all_white_drones),
            'black_boats': len(self.black_boats),
            'game_time': self.current_time,
            'drone_manager_info': {
                'type': type(self.drone_manager).__name__,
                'debug_mode': self.drone_manager.debug_mode
            }
        }
